/*
 * @Author: Sian Li lisian_magic@163.com
 * @Date: 2024-07-01 17:20:20
 * @LastEditors: Sian Li lisian_magic@163.com
 * @LastEditTime: 2024-07-01 17:58:21
 * @FilePath: /hudar/src/hudar/src/RadarClass/Detectors/include/ArmorDetector.h
 * @Description:
 */
#ifndef __ARMORDETECTOR_H
#define __ARMORDETECTOR_H

#include "../../Common/include/algorithm.h"
#include "../../TRTInference/TRTInfer/include/Inference.h"

/**
 * @brief 装甲板识别
 * 装甲板识别器，输出装甲板BBox和对应分割框
 */
class ArmorDetector {
  public:
    typedef std::shared_ptr<ArmorDetector> Ptr;

  private:
    TRTInferV1::TRTInfer armorTensorRT = TRTInferV1::TRTInfer(0);
    vector<bboxAndRect> results;

    string TensorRTEnginePath;
    string OnnxPath;

#ifdef ExperimentalOutput
    int detectedArmorNumThisFrame = 0;
    float averageConfThisFrame = 0.0;
#endif

  private:
    vector<Mat> preProcess(Mat &image, vector<DetectBox> &movingTargets);
    void reBuildBoxs(vector<vector<TRTInferV1::DetectionObj>> &armors,
                     vector<DetectBox> &boxs, vector<Mat> &img);

  public:
    ArmorDetector(string engine_path, string onnx_path);
    ~ArmorDetector();

    void accessModelTest();
    bool initModel();
    vector<bboxAndRect> infer(Mat &image, vector<DetectBox> &targets);
    void unInit();

#ifdef UseOneLayerInfer
    vector<bboxAndRect> infer(Mat &image);
#endif
};

#endif